#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/slab.h>
#include <linux/scatterlist.h>
#include <crypto/akcipher.h>
#include <linux/crypto.h>
#include <linux/string.h>
#include <linux/err.h>
#include <linux/gfp.h>

MODULE_LICENSE("GPL");
MODULE_AUTHOR("Tongyi Kernel Module");
MODULE_DESCRIPTION("Digital Signature and Verification using Crypto API in Linux Kernel 6.6");

// 模块全局变量
static struct crypto_akcipher *tfm_rsa = NULL;

// 模块初始化函数
static int __init signature_module_init(void)
{
    printk(KERN_INFO "Digital Signature Module: Initializing\n");
    
    // 初始化RSA算法 - 修正函数调用参数
    tfm_rsa = crypto_alloc_akcipher("rsa", 0, 0);
    if (IS_ERR(tfm_rsa)) {
        printk(KERN_ERR "Failed to allocate RSA transform: %ld\n", PTR_ERR(tfm_rsa));
        return PTR_ERR(tfm_rsa);
    }
    
    printk(KERN_INFO "Digital Signature Module: Successfully initialized\n");
    return 0;
}

// 模块清理函数
static void __exit signature_module_exit(void)
{
    if (tfm_rsa) {
        crypto_free_akcipher(tfm_rsa);
        printk(KERN_INFO "Digital Signature Module: Freed RSA transform\n");
    }
    
    printk(KERN_INFO "Digital Signature Module: Exited\n");
}

/**
 * 生成数字签名
 * @param data: 待签名的数据缓冲区
 * @param data_len: 数据长度
 * @param signature: 输出的签名结果缓冲区
 * @param signature_len: 签名长度（输入时为缓冲区大小，输出时为实际签名长度）
 * @param private_key: 私钥数据
 * @param key_len: 私钥长度
 * @return: 0表示成功，负数表示错误
 */
int generate_signature(const uint8_t *data, size_t data_len, 
                      uint8_t *signature, size_t *signature_len, 
                      const uint8_t *private_key, size_t key_len)
{
    struct akcipher_request *req = NULL;
    struct scatterlist src_sg, dst_sg;
    struct crypto_akcipher *local_tfm;
    void *src_buf = NULL, *dst_buf = NULL;
    int ret = 0;
    
    if (!data || !signature || !signature_len || !private_key) {
        printk(KERN_ERR "Invalid input parameters\n");
        return -EINVAL;
    }
    
    if (!tfm_rsa) {
        printk(KERN_ERR "Crypto transform not initialized\n");
        return -ENODEV;
    }
    
    local_tfm = tfm_rsa;
    
    // 分配请求结构
    req = akcipher_request_alloc(local_tfm, GFP_KERNEL);
    if (!req) {
        printk(KERN_ERR "Failed to allocate akcipher request\n");
        return -ENOMEM;
    }
    
    // 分配源数据缓冲区（连续物理内存）
    src_buf = kmalloc(data_len, GFP_KERNEL);
    if (!src_buf) {
        ret = -ENOMEM;
        goto cleanup_req;
    }
    memcpy(src_buf, data, data_len);
    
    // 分配目标签名缓冲区（连续物理内存）
    dst_buf = kmalloc(*signature_len, GFP_KERNEL);
    if (!dst_buf) {
        ret = -ENOMEM;
        goto cleanup_src;
    }
    
    // 设置scatterlist
    sg_init_one(&src_sg, src_buf, data_len);
    sg_init_one(&dst_sg, dst_buf, *signature_len);
    
    // 设置请求
    akcipher_request_set_crypt(req, &src_sg, &dst_sg, data_len, *signature_len);
    
    // 设置私钥（在实际应用中需要设置私钥）
    // 这里仅作为框架，实际需要通过crypto_akcipher_set_priv_key设置
    // 由于内核模块中直接使用私钥有安全风险，实际应用中需要从安全存储加载
    
    // 执行签名操作
    ret = crypto_akcipher_sign(req);
    if (ret) {
        if (ret == -ENOSYS) {
            printk(KERN_ERR "Sign operation not supported by algorithm\n");
            ret = -EOPNOTSUPP;
        } else {
            printk(KERN_ERR "Failed to generate signature: %d\n", ret);
        }
        goto cleanup_dst;
    }
    
    // 更新实际签名长度
    *signature_len = req->dst_len;
    
    // 复制结果到输出缓冲区
    memcpy(signature, dst_buf, req->dst_len);
    
    printk(KERN_INFO "Successfully generated signature of length %u\n", req->dst_len);

cleanup_dst:
    kfree(dst_buf);
cleanup_src:
    kfree(src_buf);
cleanup_req:
    akcipher_request_free(req);
    
    return ret;
}

/**
 * 验证数字签名
 * @param data: 原始数据缓冲区
 * @param data_len: 数据长度
 * @param signature: 签名数据
 * @param signature_len: 签名长度
 * @param public_key: 公钥数据
 * @param key_len: 公钥长度
 * @return: 0表示验证成功，负数表示错误，正数表示验证失败
 */
int verify_sign(const uint8_t *data, size_t data_len, 
                    const uint8_t *signature, size_t signature_len, 
                    const uint8_t *public_key, size_t key_len)
{
    struct akcipher_request *req = NULL;
    struct scatterlist src_sg, dst_sg;
    struct crypto_akcipher *local_tfm;
    void *src_buf = NULL;
    int ret = 0;
    
    if (!data || !signature || !public_key) {
        printk(KERN_ERR "Invalid input parameters\n");
        return -EINVAL;
    }
    
    if (!tfm_rsa) {
        printk(KERN_ERR "Crypto transform not initialized\n");
        return -ENODEV;
    }
    
    local_tfm = tfm_rsa;
    
    // 分配请求结构
    req = akcipher_request_alloc(local_tfm, GFP_KERNEL);
    if (!req) {
        printk(KERN_ERR "Failed to allocate akcipher request\n");
        return -ENOMEM;
    }
    
    // 分配源数据缓冲区（签名数据）
    src_buf = kmalloc(signature_len, GFP_KERNEL);
    if (!src_buf) {
        ret = -ENOMEM;
        goto cleanup_req;
    }
    memcpy(src_buf, signature, signature_len);
    
    // 分配目标缓冲区用于验证操作（使用数据长度作为缓冲区大小）
    uint8_t *dst_buf = kmalloc(data_len, GFP_KERNEL);
    if (!dst_buf) {
        ret = -ENOMEM;
        goto cleanup_src;
    }
    
    // 设置scatterlist
    sg_init_one(&src_sg, src_buf, signature_len);
    sg_init_one(&dst_sg, dst_buf, data_len);
    
    // 设置请求
    akcipher_request_set_crypt(req, &src_sg, &dst_sg, signature_len, data_len);
    
    // 执行验证操作
    ret = crypto_akcipher_verify(req);
    if (ret == -ENOSYS) {
        // 如果验证操作不支持，直接返回错误
        printk(KERN_ERR "Verify operation not supported by algorithm\n");
        ret = -EOPNOTSUPP;
        goto cleanup_dst;
    } else if (ret) {
        printk(KERN_ERR "Signature verification failed: %d\n", ret);
        goto cleanup_dst;
    } else {
        printk(KERN_INFO "Signature verification successful\n");
    }

cleanup_dst:
    kfree(dst_buf);
cleanup_src:
    kfree(src_buf);
cleanup_req:
    akcipher_request_free(req);
    
    return ret;
}

// 导出函数供其他模块使用
EXPORT_SYMBOL(generate_signature);
EXPORT_SYMBOL(verify_sign);

module_init(signature_module_init);
module_exit(signature_module_exit);